package org.weed.gadgets;

import java.io.BufferedReader;
import java.io.File;
import java.io.InputStreamReader;
import java.nio.file.FileSystems;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.MessageDigest;

import javax.crypto.Cipher;
import javax.crypto.spec.SecretKeySpec;

public class AESCipher {
	private final String HASH_FUNCTION = "MD5";	 //SHA-256 for 256-bit key (AES not supports 256-bit key)
	private Cipher cipher;
	private SecretKeySpec keySpec;
	
	public AESCipher(String password) {
		try {
			keySpec = new SecretKeySpec(hash(password), "AES");
			cipher = Cipher.getInstance("AES");
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	public byte[] encrypt(byte[] bytes) {
		try {
			cipher.init(Cipher.ENCRYPT_MODE, keySpec);
			return cipher.doFinal(bytes);
		} catch (Exception e) { 
			e.printStackTrace(); 
			return null;
		}
	}
	
	public boolean encrypt(String path) {
		path = path.replace("\\", "\\\\");
		try {
			Path file = FileSystems.getDefault().getPath(path);
			byte[] text = Files.readAllBytes(file);
			Files.write(file, encrypt(text));
			return true;
		} catch (Exception e) {
			//e.printStackTrace();
			return false;
		}
	}
	
	public boolean decrypt(String path) {
		path = path.replace("\\", "\\\\");
		try {
			Path file = FileSystems.getDefault().getPath(path);
			byte[] secret = Files.readAllBytes(file);
			Files.write(file, decrypt(secret));
			return true;
		} catch (Exception e) {
			//e.printStackTrace();
			return false;
		}
	}
	
	public byte[] decrypt(byte[] bytes) {
		try {
			cipher.init(Cipher.DECRYPT_MODE, keySpec);
			return cipher.doFinal(bytes);
		} catch (Exception e) {
			e.printStackTrace();
			return null;
		}
	}
	
	public byte[] hash(String msg) {
		try {
			MessageDigest hasher = MessageDigest.getInstance(HASH_FUNCTION);
			return hasher.digest(msg.getBytes());
		} catch (Exception e) {
			e.printStackTrace();
			return null;
		}
	}
	
	public static void main(String[] args) {
		try {
			BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
			AESCipher cipher = null;
			String password = "";
			while (password.length() <= 0) {
				System.out.print("Please setup a password: ");
				password = reader.readLine().trim();
				cipher = new AESCipher(password);
			}
			System.out.println("Type hint for hint");
			String command = null;
			String[] commands = null;
			while(true) {
				System.out.print(">");
				command = reader.readLine().trim();
				if ("HINT".equals(command.toUpperCase())) {
					System.out.println("To encrpt or decrypt a file, type \"-e|-d {filepath}\"");
					System.out.println("To reset password, type \"reset {password}\"");
					System.out.println("To quit the program, click \"x\"");
				} else {
					commands = command.split(" ", 2);
					try {
						String command0 = commands[0].toUpperCase();
						String command1 = commands[1];
						if (command1.length() > 0) {
							if ("RESET".equals(command0)) {
								cipher = new AESCipher(command1);
								System.out.println("Password changed.");
							} else if ("-E".equals(command0)) {
								if (cipher.encrypt(command1)) {
									System.out.println("File encrypted.");
								} else {
									System.out.println("Fail to encrypt the specified file.");
								}
							} else if ("-D".equals(command0)) {
								if (cipher.decrypt(command1)) {
									System.out.println("File decrypted.");
								} else {
									System.out.println("Fail to decrypt the specified file.");
								}
							}
						} else {
							throw new IllegalArgumentException("Invalid command");
						}
					} catch (Exception e) {
						System.out.println("Invalid command, type hint for hint.");
					}
				}			
			}
		} catch (Exception e) {
			
		}
	}
}
